#!/usr/bin/env python3
"""
快速设置竹马法考token
"""
import sys
import os

# 添加项目路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from app.scrapers.zhuma_scraper import zhuma_scraper
from app.core.database import RedisClient
import json
from datetime import datetime

def set_token(token: str, username: str = "default_user"):
    """
    设置竹马法考token
    
    Args:
        token: token值
        username: 用户名（可选）
    """
    print(f"正在设置token...")
    
    # 设置token（token、stoken、mtoken通常相同）
    zhuma_scraper.username = username
    zhuma_scraper.token = token
    zhuma_scraper.stoken = token
    zhuma_scraper.mtoken = token
    zhuma_scraper.logged_in = True
    
    # 更新请求头
    zhuma_scraper._update_headers()
    
    # 保存到Redis
    try:
        redis = RedisClient.get_client()
        if redis:
            session_key = f"zhuma:session:{username}"
            session_data = {
                "username": username,
                "cookies": {},  # token方式不需要cookies
                "token": token,
                "stoken": token,
                "mtoken": token,
                "logged_in_at": datetime.now().isoformat()
            }
            # 保存7天
            redis.setex(session_key, 7 * 24 * 3600, json.dumps(session_data, ensure_ascii=False))
            print(f"✅ Token已保存到Redis")
    except Exception as e:
        print(f"⚠️ 保存到Redis失败（将使用内存存储）: {e}")
    
    print(f"✅ Token设置成功")
    print(f"   用户名: {username}")
    print(f"   Token: {token[:50]}...")
    
    # 测试登录状态
    status = zhuma_scraper.check_login_status()
    if status.get("logged_in"):
        print(f"✅ 登录状态验证成功")
    else:
        print(f"⚠️ 登录状态验证失败，可能需要完整的cookies")


if __name__ == "__main__":
    # 从命令行参数获取token
    if len(sys.argv) < 2:
        print("用法: python set_zhuma_token.py <token> [username]")
        print("\n示例:")
        print("  python set_zhuma_token.py c7a588bbcaea1470f1b60437dbdbe806d8c8f1a28775efcbfddd41498d684644d4b8d934ace488c5c4426cc496282fc1144954df3c72dcddc5d7718a5bbf7a2ce826793907f21c6ddb8e0ef957922b0e98a417a7972a593f032e076c4defe2f79c370a690ba1a348ecfecde8c88bebbed96a9830393374fc236646b6a219545f")
        sys.exit(1)
    
    token = sys.argv[1]
    username = sys.argv[2] if len(sys.argv) > 2 else "default_user"
    
    set_token(token, username)

